import os
import requests
import m3u8
import shutil
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import urljoin

# 配置
OUTPUT_VIDEO_NAME = "xxx.mp4"  # 输出视频名称
TS_DIR = "ts_files"                # ts 文件临时存储目录
MAX_WORKERS = 10                   # 并发线程数

def download_file(args):
    url, filename = args
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            with open(filename, 'wb') as f:
                f.write(response.content)
            return True
        else:
            print(f"[错误] 下载失败: {url} 状态码: {response.status_code}")
            return False
    except Exception as e:
        print(f"[异常] 下载失败: {url} 异常: {e}")
        return False

def download_ts_files(base_url, ts_list):
    if not os.path.exists(TS_DIR):
        os.makedirs(TS_DIR)

    tasks = []
    for i, segment in enumerate(ts_list.segments):
        ts_url = urljoin(base_url, segment.uri)
        ts_file_name = os.path.join(TS_DIR, f"{i:05d}.ts")
        tasks.append((ts_url, ts_file_name))

    print(f"开始下载 {len(tasks)} 个 ts 文件...")
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        results = executor.map(download_file, tasks)

    success = sum(1 for r in results if r)
    print(f"成功下载 {success}/{len(tasks)} 个 ts 文件。")

def merge_ts_files():
    print("开始合并 ts 文件...")
    with open(OUTPUT_VIDEO_NAME, 'wb') as merged:
        for i in range(len(os.listdir(TS_DIR))):
            ts_file = os.path.join(TS_DIR, f"{i:05d}.ts")
            with open(ts_file, 'rb') as seg:
                merged.write(seg.read())
    print(f"视频已合并为: {OUTPUT_VIDEO_NAME}")

def clean_up():
    print(f"正在清理临时文件夹: {TS_DIR}")
    shutil.rmtree(TS_DIR)

def main(m3u8_url):
    response = requests.get(m3u8_url)
    if response.status_code != 200:
        print("无法获取 m3u8 文件。")
        return

    m3u8_obj = m3u8.loads(response.text)
    base_url = m3u8_url.rsplit('/', 1)[0] + '/'

    if m3u8_obj.is_variant:
        print("检测到是多清晰度流，请提供具体分辨率的 m3u8 链接。")
        return

    download_ts_files(base_url, m3u8_obj)
    merge_ts_files()
    clean_up()
    print("✅ 下载完成！")

if __name__ == '__main__':
    m3u8_url = input("请输入视频的 m3u8 文件链接：").strip()
    main(m3u8_url)